/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia;



import com.bff.gaia.unified.runners.core.metrics.MetricsPusher;

import com.bff.gaia.unified.sdk.Pipeline;

import com.bff.gaia.unified.sdk.PipelineResult;

import com.bff.gaia.unified.sdk.PipelineRunner;

import com.bff.gaia.unified.sdk.metrics.MetricsEnvironment;

import com.bff.gaia.unified.sdk.metrics.MetricsOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptionsValidator;

import com.bff.gaia.unified.sdk.runners.TransformHierarchy;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.View;

import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import com.bff.gaia.unified.vendor.guava.com.google.common.base.Joiner;

import com.bff.gaia.api.common.JobExecutionResult;

import com.bff.gaia.client.program.DetachedEnvironment;

import com.bff.gaia.runtime.jobgraph.JobGraph;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.util.*;



import static com.bff.gaia.unified.runners.core.construction.PipelineResources.detectClassPathResourcesToStage;



/**

 * A {@link PipelineRunner} that executes the operations in the pipeline by first translating them

 * to a Gaia Plan and then executing them either locally or on a Gaia cluster, depending on the

 * configuration.

 */

public class GaiaRunner extends PipelineRunner<PipelineResult> {



  private static final Logger LOG = LoggerFactory.getLogger(GaiaRunner.class);



  /** Provided options. */

  private final GaiaPipelineOptions options;



  /**

   * Construct a runner from the provided options.

   *

   * @param options Properties which configure the runner.

   * @return The newly created runner.

   */

  public static GaiaRunner fromOptions(PipelineOptions options) {

    GaiaPipelineOptions gaiaOptions =

        PipelineOptionsValidator.validate(GaiaPipelineOptions.class, options);

    ArrayList<String> missing = new ArrayList<>();



    if (gaiaOptions.getAppName() == null) {

      missing.add("appName");

    }

    if (missing.size() > 0) {

      throw new IllegalArgumentException(

          "Missing required values: " + Joiner.on(',').join(missing));

    }



    if (gaiaOptions.getFilesToStage() == null) {

      gaiaOptions.setFilesToStage(

          detectClassPathResourcesToStage(GaiaRunner.class.getClassLoader()));

      LOG.info(

          "PipelineOptions.filesToStage was not specified. "

              + "Defaulting to files from the classpath: will stage {} files. "

              + "Enable logging at DEBUG level to see which files will be staged.",

          gaiaOptions.getFilesToStage().size());

      LOG.debug("Classpath elements: {}", gaiaOptions.getFilesToStage());

    }



    return new GaiaRunner(gaiaOptions);

  }



  private GaiaRunner(GaiaPipelineOptions options) {

    this.options = options;

    this.ptransformViewsWithNonDeterministicKeyCoders = new HashSet<>();

  }



  @Override

  public PipelineResult run(Pipeline pipeline) {

    logWarningIfPCollectionViewHasNonDeterministicKeyCoder(pipeline);



    MetricsEnvironment.setMetricsSupported(true);



    LOG.info("Executing pipeline using GaiaRunner.");



    GaiaPipelineExecutionEnvironment env = new GaiaPipelineExecutionEnvironment(options);



    LOG.info("Translating pipeline to Gaia program.");

    env.translate(pipeline);



    JobExecutionResult result;

    try {

      LOG.info("Starting execution of Gaia program.");

      result = env.executePipeline();

    } catch (Exception e) {

      LOG.error("Pipeline execution failed", e);

      throw new RuntimeException("Pipeline execution failed", e);

    }

    return createPipelineResult(result, options);

  }



  static PipelineResult createPipelineResult(JobExecutionResult result, PipelineOptions options) {

    if (result instanceof DetachedEnvironment.DetachedJobExecutionResult) {

      LOG.info("Pipeline submitted in Detached mode");

      // no metricsPusher because metrics are not supported in detached mode

      return new GaiaDetachedRunnerResult();

    } else {

      LOG.info("Execution finished in {} msecs", result.getNetRuntime());

      Map<String, Object> accumulators = result.getAllAccumulatorResults();

      if (accumulators != null && !accumulators.isEmpty()) {

        LOG.info("Final accumulator values:");

        for (Map.Entry<String, Object> entry : result.getAllAccumulatorResults().entrySet()) {

          LOG.info("{} : {}", entry.getKey(), entry.getValue());

        }

      }

      GaiaRunnerResult gaiaRunnerResult =

          new GaiaRunnerResult(accumulators, result.getNetRuntime());

      MetricsPusher metricsPusher =

          new MetricsPusher(

              gaiaRunnerResult.getMetricsContainerStepMap(),

              options.as(MetricsOptions.class),

              gaiaRunnerResult);

      metricsPusher.start();

      return gaiaRunnerResult;

    }

  }



  /** For testing. */

  @VisibleForTesting

  public GaiaPipelineOptions getPipelineOptions() {

    return options;

  }



  @Override

  public String toString() {

    return "GaiaRunner#" + hashCode();

  }



  /** A set of {@link View}s with non-deterministic key coders. */

  Set<PTransform<?, ?>> ptransformViewsWithNonDeterministicKeyCoders;



  /** Records that the {@link PTransform} requires a deterministic key coder. */

  void recordViewUsesNonDeterministicKeyCoder(PTransform<?, ?> ptransform) {

    ptransformViewsWithNonDeterministicKeyCoders.add(ptransform);

  }



  /** Outputs a warning about PCollection views without deterministic key coders. */

  private void logWarningIfPCollectionViewHasNonDeterministicKeyCoder(Pipeline pipeline) {

    // We need to wait till this point to determine the names of the transforms since only

    // at this time do we know the hierarchy of the transforms otherwise we could

    // have just recorded the full names during apply time.

    if (!ptransformViewsWithNonDeterministicKeyCoders.isEmpty()) {

      final SortedSet<String> ptransformViewNamesWithNonDeterministicKeyCoders = new TreeSet<>();

      pipeline.traverseTopologically(

          new Pipeline.PipelineVisitor.Defaults() {



            @Override

            public void visitPrimitiveTransform(TransformHierarchy.Node node) {

              if (ptransformViewsWithNonDeterministicKeyCoders.contains(node.getTransform())) {

                ptransformViewNamesWithNonDeterministicKeyCoders.add(node.getFullName());

              }

            }



            @Override

            public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {

              if (ptransformViewsWithNonDeterministicKeyCoders.contains(node.getTransform())) {

                ptransformViewNamesWithNonDeterministicKeyCoders.add(node.getFullName());

              }

              return CompositeBehavior.ENTER_TRANSFORM;

            }

          });



      LOG.warn(

          "Unable to use indexed implementation for View.AsMap and View.AsMultimap for {} "

              + "because the key coder is not deterministic. Falling back to singleton "

              + "implementation which may cause memory and/or performance problems. Future major "

              + "versions of the Gaia runner will require deterministic key coders.",

          ptransformViewNamesWithNonDeterministicKeyCoders);

    }

  }



  @VisibleForTesting

  JobGraph getJobGraph(Pipeline p) {

    GaiaPipelineExecutionEnvironment env = new GaiaPipelineExecutionEnvironment(options);

    return env.getJobGraph(p);

  }

}